import argparse  
import os  
import shutil
import subprocess  

# 遍历目录 并获取以特定后缀结尾的文件
def get_files_by_extension(directory, extensions):  
    files = os.listdir(directory)  
    result = {}  
    for file in files:  
        if file.endswith(tuple(extensions)):  
            base_name = os.path.splitext(file)[0]  
            if base_name not in result:  
                result[base_name] = []  
            result[base_name].append(file)  
    return result 

# 生成【删除表空间】语句
def discard_table_space(table_name):  
    # 设置文件路径和名称  
    file_path = './dist/discard_table_space.sql'
    # 检查目标文件夹是否存在，如果不存在则创建它  
    if not os.path.exists('./dist'):  
        os.makedirs('dist')  
    
    # 检查文件是否存在  
    if not os.path.exists(file_path):  
        # 如果文件不存在，使用'w'模式创建它  
        with open(file_path, 'w') as f:  
            f.write(f"alter table `{table_name}` discard tablespace;\n")  
    else:  
        # 如果文件存在，使用'a'模式打开它并追加内容  
        with open(file_path, 'a') as f:  
            f.write(f"alter table `{table_name}` discard tablespace;\n")        

# 生成【导入表空间】语句
def import_table_space(table_name):  
    # 设置文件路径和名称  
    file_path = './dist/import_table_space.sql'
    # 检查目标文件夹是否存在，如果不存在则创建它  
    if not os.path.exists('./dist'):  
        os.makedirs('dist')  
    
    # 检查文件是否存在  
    if not os.path.exists(file_path):  
        # 如果文件不存在，使用'w'模式创建它  
        with open(file_path, 'w') as f:  
            f.write(f"alter table `{table_name}` import tablespace;\n")  
    else:  
        # 如果文件存在，使用'a'模式打开它并追加内容  
        with open(file_path, 'a') as f:  
            f.write(f"alter table `{table_name}` import tablespace;\n")        

# 拷贝文件到指定目录
def copy_2_mysql_data(file_list, dest_folder, source_dir = './data'):
    # 检查目标文件夹是否存在，如果不存在则创建它  
    if not os.path.exists(dest_folder):  
        os.makedirs(dest_folder)  
    
    # 遍历源文件列表，逐个复制到目标文件夹  
    for file in file_list:  
        # 获取源文件路径  
        src_path = source_dir+'/'+file  
    
        # 获取目标文件路径  
        dst_path = os.path.join(dest_folder, os.path.basename(file))  
    
        # 复制文件  
        shutil.copy2(src_path, dst_path)      
  
def main():          
    parser = argparse.ArgumentParser(description="批量恢复数据库")  
  
    # 添加必需的参数  
    parser.add_argument('--source_dir', type=str, required=True, help="待恢复文件目录")   
      
    # 可选的参数  
    parser.add_argument('--table', type=str, help="指定要恢复的表 不填默认为目录下全部")  
  
    args = parser.parse_args()  
  
    # 验证路径是否存在  
    if not os.path.exists(args.source_dir):  
        print("错误：源目录不存在！")  
        exit(1)  
  
    print("源目录: ", args.source_dir)  
    if args.table is not None:  
        print("要恢复的表: ", args.table)  
    else:  
        print("未指定要恢复的表。即将恢复所有")  

    # 获取特定扩展名的文件  
    files = get_files_by_extension(args.source_dir, [".frm", ".ibd", ".MYI", ".MYD"])  
    if args.table is not None: 
        if args.table not in files:
            print("指定恢复的表文件不存在")

    # 删除指定文件
    if os.path.exists('./dist/discard_table_space.sql'):  
        os.remove('./dist/discard_table_space.sql')
    if os.path.exists('./dist/import_table_space.sql'):      
        os.remove('./dist/import_table_space.sql')
    if os.path.exists('./dist/create_table.sql'):  
        os.remove('./dist/create_table.sql')

    print("文件遍历中……")
    # 循环遍历files
    for key, value in files.items():
        if args.table is not None and  key != args.table:
            continue 
        # 循环遍历value列表中 找到同时包含*.frm和*.ibd的数据 即为innodb的表相关
        if any(file.endswith('.ibd') for file in value) and any(file.endswith('.frm') for file in value):  
            # 生成删除表空间语句
            discard_table_space(key)
            # 生成导入表空间语句
            import_table_space(key)
            # 将innodb的frm和ibd文件拷贝出来
            copy_2_mysql_data([key+'.frm'], './dist/frm', args.source_dir)
            copy_2_mysql_data([key+'.ibd'], './dist/ibd', args.source_dir)
        elif any(file.endswith('.MYD') for file in value) and any(file.endswith('.MYI') for file in value) and any(file.endswith('.frm') for file in value):
            # myisam的直接拷贝到单独目录，后期直接放到数据文件目录中即可恢复
            copy_2_mysql_data(value, './dist/myisam', args.source_dir)

    print("========================执行结果==============================")
    print("删除表空间语句生成成功，请查看dist/discard_table_space.sql")
    print("导入表空间语句生成成功，请查看dist/import_table_space.sql")

    # 使用dbsake提取frm的结构到sql文件 需要赋予dbsake可执行权限
    subprocess.run(['chmod', 'u+x', 'dbsake'])
    # 执行语句 ./dbsake frmdump ./dist/frm/*.frm >> ./dist/create_table.sql  
    # 第一步：运行 ./dbsake frmdump ./dist/frm/*.frm
    process1 = subprocess.Popen(['./dbsake', 'frmdump', './dist/frm/*.frm'], stdout=subprocess.PIPE)  
  
    # 第二步：将输出附加到 ./dist/create_table.sql 文件  
    with open('./dist/create_table.sql', 'a') as file:  
        file.write(process1.stdout.read().decode('utf-8'))

    print("数据库Innodb表结构已获取成功，请查看dist/create_table.sql")
    print("========================后续操作==============================")
    print("1. 【创建表】在新的数据库中执行dist/create_table.sql中的语句")
    print("2. 【删除表空间】在新的数据库中执行dist/discard_table_space.sql中的语句")
    print("3. 【拷贝表的源数据】将dist/ibd目录和dist/myisam目录下的所有文件都拷贝到mysql的数据库目录下")
    print("4. 【导入表空间】在新的数据库中执行dist/import_table_space.sql中的语句")
    print("操作完成，请检查数据库是否恢复成功")

  
if __name__ == "__main__":  
    main()